import json
import time
import traceback
import cv2
import numpy as np
from fastapi import FastAPI, HTTPException
from paddleocr import PaddleOCR
from pydantic import BaseModel
import os
import shutil
from datetime import datetime
import base64

app = FastAPI()
ocr = PaddleOCR(
    text_detection_model_name="PP-OCRv5_mobile_det",
    text_recognition_model_name="PP-OCRv5_mobile_rec",
    use_doc_orientation_classify=False,
    use_doc_unwarping=False,
    use_textline_orientation=False,
    device="cpu",
)

# 定义请求模型：接收图像数据列表
class OCRRequest(BaseModel):
    """
    定义 OCR 请求的数据模型，包含图像数据列表。
    该列表由 numpy.ndarray 转换而来。
    """
    image_data: str  # 图像数据列表（由 numpy.ndarray 转换而来）


def list_to_numpy(image_list):
    """
    将列表转换为 numpy.ndarray。
    :param image_list: 图像数据列表
    :return: 转换后的 numpy.ndarray 数组
    :raises ValueError: 当图像数据为空或维度不符合要求时抛出异常
    """
    if not image_list:
        raise ValueError("图像数据为空")
    # 根据列表结构推断数据类型（通常为 uint8）
    img_np = np.array(image_list, dtype=np.uint8)
    # 检查维度（确保为 HWC 格式，如 [height, width, channels]）
    if len(img_np.shape) != 3:
        raise ValueError("图像维度错误，需为 HWC 格式")
    return img_np


@app.post("/ocr")
async def ocr_endpoint(request: OCRRequest):
    """
    处理 OCR 请求的端点
    """
    try:
        start_time = time.time()
        

        img_data = base64.b64decode(request.image_data)
        img_np = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_COLOR)

        if img_np is None:
            raise ValueError("无法解码图像数据")
       
        # 将列表转换为 numpy.ndarray
        # img_np = list_to_numpy(request.image_data)

        # 生成基于时间的唯一标识符
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
        save_dir = f"output/{timestamp}"
        os.makedirs(save_dir, exist_ok=True)
        
        # 指定JSON文件路径
        json_file_path = f"{save_dir}_res.json"

        ocr_start = time.time()
        # 执行 OCR 识别（直接传入 numpy.ndarray）
        results = ocr.predict_iter(img_np)

        # 打印和保存可视化结果
        for res in results:
            #res.print()
            res.save_to_json(json_file_path)
            #res.save_to_img(save_path=f"{save_dir}_res.png")

        # 读取 JSON 文件内容
        with open(json_file_path, 'r', encoding='utf-8') as f:
            ocr_result = json.load(f)

        # 删除文件夹
        # shutil.rmtree(save_dir)

        ocr_elapsed = time.time() - ocr_start  # OCR 服务处理耗时
        total_elapsed = time.time() - start_time  # 从调用函数到获取结果总耗时
        
        print(f"OCR 识别成功 | 服务处理耗时: {ocr_elapsed:.4f} 秒 | 总耗时: {total_elapsed:.4f} 秒")

        return {
            "status": "success",
            "original_shape": str(img_np.shape),
            "ocr_result": ocr_result,
            "processing_time": timestamp  # 返回处理时间戳
        }

    except Exception as e:
        traceback.print_exc()
        # 明确指定原始异常
        raise HTTPException(status_code=500, detail=f"OCR 处理失败: {str(e)}") from e


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        # 改为导入字符串："模块名:应用实例名"
        "ocr_http_server:app",  # 假设文件名为 match_http_server.py
        host="0.0.0.0",
        port=8003,
        reload=False,  # 开发模式启用热重载
        # workers=4,  # 生产环境按需开启
    )